package com.ks.test.app;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.Objects;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;

public class MyTransformer implements ClassFileTransformer {


    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {

        //跳过java自带方法
        if (className.startsWith("java") || className.startsWith("sun")) {
            return classfileBuffer;
        }

        //好像使用premain这个className是没问题的，但使用attach时className的.变成了/，所以如果是attach，那么这里需要替换
        className = className.replace("/", ".");

        //只处理MyApp类
        if (!className.endsWith("MyApp")) {
            return classfileBuffer;
        }

        try {
            ClassPool classPool = ClassPool.getDefault();
            CtClass ctClass = classPool.get(className);
            CtMethod[] declaredMethods = ctClass.getDeclaredMethods();

            for (CtMethod declaredMethod : declaredMethods) {
                //只处理printSth方法
                if (Objects.equals("printHello", declaredMethod.getName())) {
                    //在方法执行前插入打印语句
                    declaredMethod.insertBefore("System.out.println(\"=====start=====\");");
                    //在方法执行后插入打印语句
                    declaredMethod.insertAfter("System.out.println(\"=====end=====\");");

                    break;
                }
            }

            return ctClass.toBytecode();

        } catch (Exception e) {
            e.printStackTrace();
        }

        return classfileBuffer;
    }
}
